#include "global_planner/global_planner.h"
#include "transform_utils/type_transform.h"
#include "transform_utils/display.h"
#include <fstream>
#include <rosbag/bag.h>
#include <pcl_conversions/pcl_conversions.h>

namespace Planner
{
    GlobalPlanner::GlobalPlanner(const ros::NodeHandle &node_handle)
        : nh_(node_handle)
    {
        init();
    }

    void GlobalPlanner::init()
    {
        nh_.param("global_planner/traj_topic", traj_topic_, string(""));
        nh_.param("global_planner/opt_seg_num", opt_seg_num_, 8);

        ompl_planner_.setParams(nh_);
        traj_opt_.setParams(nh_);

        int bspline_order;
        double des_gap;
        nh_.param("global_planner/bspline_order", bspline_order, 3);
        nh_.param("global_planner/des_gap", des_gap, 1.5);
        bspline_smooth_.setParams(bspline_order, des_gap);

        // 读取点云地图
        cloud_map_ = boost::make_shared<pcl::PointCloud<pcl::PointXYZ>>(pcl::PointCloud<pcl::PointXYZ>());
        kdtree_ = boost::make_shared<pcl::KdTreeFLANN<pcl::PointXYZ>>(pcl::KdTreeFLANN<pcl::PointXYZ>());

        string cloud_file;
        nh_.param("global_planner/cloud_file", cloud_file, string(""));
        if (pcl::io::loadPCDFile(cloud_file, *cloud_map_) == -1)
        {
            ROS_ERROR("Cannot read pcd file");
            status_ = INIT_FAIL;
        }
        else
        {
            kdtree_->setInputCloud(cloud_map_);
            ompl_planner_.setCloudMap(cloud_map_, kdtree_);
            traj_opt_.setCloudMap(cloud_map_, kdtree_);
            status_ = INIT_SUCCESS;
        }

        string waypoint_topic;
        nh_.param("global_planner/waypoint_topic", waypoint_topic, std::string(""));
        waypoint_sub_ = nh_.subscribe(waypoint_topic, 1, &GlobalPlanner::waypointCallback, this);
        traj_process_server_ = nh_.advertiseService("traj_process", &GlobalPlanner::trajProcessCallback, this);

        // 用于可视化
        waypoints_pub_ = nh_.advertise<visualization_msgs::MarkerArray>("waypoints_vis", 1);
        origin_path_pub_ = nh_.advertise<nav_msgs::Path>("origin_path_vis", 1);
        smooth_traj_pub_ = nh_.advertise<nav_msgs::Path>("smooth_traj_vis", 1);
        sample_traj_pub_ = nh_.advertise<nav_msgs::Path>("sample_traj_vis", 1);
        optimized_traj_pub_ = nh_.advertise<nav_msgs::Path>("optimized_traj_vis", 1);
        cloud_map_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("cloud_map_vis", 5);

        // NOTE 只发送一次接受不到数据，不知道为啥
        sensor_msgs::PointCloud2 cloud_map_msgs;
        pcl::toROSMsg(*cloud_map_, cloud_map_msgs);
        cloud_map_msgs.header.frame_id = "world";
        for (int i = 0; i < 2; i++)
        {
            cloud_map_pub_.publish(cloud_map_msgs);
            ros::Duration(1.0).sleep();
        }
        ROS_INFO("Global planner initialized");
    }

    void GlobalPlanner::waypointCallback(const geometry_msgs::PoseStamped &msg)
    {
        cout << msg.pose.position.x << " " << msg.pose.position.y << " " << msg.pose.position.z << endl;
        switch (status_)
        {
        case INIT_FAIL:
        {
            ROS_WARN("Initialization failed");
            break;
        }

        case INIT_SUCCESS:
        {
            Vector3d pt;
            pt[0] = msg.pose.position.x;
            pt[1] = msg.pose.position.y;
            pt[2] = msg.pose.position.z;
            if (ompl_planner_.setStart(pt))
            {
                start_pos_ = pt;
                waypoints_.push_back(start_pos_);
                ROS_INFO("Set start position to (%f, %f, %f)", start_pos_[0], start_pos_[1], start_pos_[2]);
                status_ = SET_START;

                displayWaypoints(waypoints_, waypoints_pub_);
            }
            else
                ROS_WARN("Position (%f, %f, %f) is unsafe", start_pos_[0], start_pos_[1], start_pos_[2]);
            break;
        }

        case SET_START:
        {
            Vector3d pt;
            pt[0] = msg.pose.position.x;
            pt[1] = msg.pose.position.y;
            pt[2] = msg.pose.position.z;
            if (ompl_planner_.setGoal(pt))
            {
                ROS_INFO("Set goal position to (%f, %f, %f)", pt[0], pt[1], pt[2]);

                // 开始路径规划
                if (ompl_planner_.plan())
                {
                    goal_pos_ = pt;
                    waypoints_.push_back(goal_pos_);
                    vector<Vector3d> path = ompl_planner_.getPath(false);
                    origin_path_.insert(origin_path_.end(), path.begin(), path.end());

                    status_ = SET_GOAL;
                    ops_ = ADD_WAYPOINTS;
                    ROS_INFO("Find a feasible path");

                    displayWaypoints(waypoints_, waypoints_pub_);
                    displayTrajectory(origin_path_, origin_path_pub_);
                }
                else
                    ROS_WARN("Cannot plan a feasible path, please choose another goal position");
            }
            else
                ROS_INFO("Position (%f, %f, %f) is unsafe", pt[0], pt[1], pt[2]);
            break;
        }

        case SET_GOAL:
        {
            Vector3d pt;
            pt[0] = msg.pose.position.x;
            pt[1] = msg.pose.position.y;
            pt[2] = msg.pose.position.z;
            if (ompl_planner_.setGoal(pt))
            {
                ompl_planner_.setStart(goal_pos_);
                ROS_INFO("Set start position to (%f, %f, %f)", goal_pos_[0], goal_pos_[1], goal_pos_[2]);
                ROS_INFO("Set goal position to (%f, %f, %f)", pt[0], pt[1], pt[2]);

                // 开始路径规划
                if (ompl_planner_.plan())
                {
                    start_pos_ = goal_pos_;
                    goal_pos_ = pt;
                    waypoints_.push_back(goal_pos_);

                    vector<Vector3d> path = ompl_planner_.getPath(false);
                    origin_path_.insert(origin_path_.end(), path.begin(), path.end());

                    ops_ = ADD_WAYPOINTS;
                    ROS_INFO("Find a feasible path");

                    displayWaypoints(waypoints_, waypoints_pub_);
                    displayTrajectory(origin_path_, origin_path_pub_);
                }
                else
                    ROS_WARN("Cannot plan a feasible path, please choose another goal position");
            }
            else
                ROS_INFO("Position (%f, %f, %f) is unsafe", goal_pos_[0], goal_pos_[1], goal_pos_[2]);
            break;
        }

        default:
            break;
        }
    }

    bool GlobalPlanner::trajProcessCallback(quad_msgs::GlobalTrajProcess::Request &request,
                                            quad_msgs::GlobalTrajProcess::Response &response)
    {
        if (origin_path_.empty())
        {
            ROS_WARN("Empty path, cannot do operations");
            response.success = false;
            return true;
        }

        switch (request.request_type)
        {
        case quad_msgs::GlobalTrajProcess::Request::SAVE_INIT_WAYPOINTS:
        {
            std::ofstream outfile;
            outfile.open(request.file_name, std::ios::out | std::ios::trunc);

            for (auto pt : origin_path_)
                outfile << pt[0] << " " << pt[1] << " " << pt[2] << "\n";
            ROS_INFO("Save origin path into file %s", request.file_name);
            break;
        }

        case quad_msgs::GlobalTrajProcess::Request::SAVE_SAMPLE_WAYPOINTS:
        {
            if (ops_ == ADD_WAYPOINTS)
                samplePath();

            std::ofstream outfile;
            outfile.open(request.file_name, std::ios::out | std::ios::trunc);

            for (auto pt : sample_path_)
                outfile << pt[0] << " " << pt[1] << " " << pt[2] << "\n";
            ROS_INFO("Save origin path into file %s", request.file_name);
            break;
        }

        // 与前两种情况保存为txt不同，该类将轨迹信息保存为rosbag
        case quad_msgs::GlobalTrajProcess::Request::SAVE_OPTIMIZE_TRAJECTORY:
        {
            if (ops_ == ADD_WAYPOINTS)
                samplePath();
            if (ops_ == SAMPLE_WAYPOINTS)
                optimizeTraj();
            if (ops_ == OPTIMIZE_FAIL)
            {
                ROS_ERROR("Error occured while optimizing trajectory");
                response.success = false;
                return true;
            }

            quad_msgs::Trajectory traj_msgs;
            coeffsToGlobalTrajMsg(coeffs_, time_seq_, traj_msgs);

            rosbag::Bag bag;
            bag.open(request.file_name, rosbag::bagmode::Write);
            bag.write(traj_topic_, ros::Time::now(), traj_msgs);

            ROS_INFO("Save optimized trajectory into file %s", request.file_name);
            response.success = true;
            break;
        }

        default:
            break;
        }

        return true;
    }

    void GlobalPlanner::samplePath()
    {
        bspline_smooth_.smoothPath(origin_path_, smooth_path_);
        bspline_smooth_.samplePath(smooth_path_, sample_path_);
        ops_ = SAMPLE_WAYPOINTS;

        displayTrajectory(smooth_path_, smooth_traj_pub_);
        displayTrajectory(sample_path_, sample_traj_pub_);
    }

    void GlobalPlanner::optimizeTraj()
    {
        int opt_times = sample_path_.size() / opt_seg_num_;
        if (opt_times == 0)
            opt_times = 1;
        // NOTE 这里轨迹阶数先固定为5阶
        coeffs_ = vector<MatrixXd>(3, MatrixXd::Zero(sample_path_.size() - 1, 6));
        time_seq_.clear();

        for (int s = 0; s < opt_times; s++)
        {
            int start_index = s * opt_seg_num_;
            int end_index = (s + 1) * opt_seg_num_;
            if (s == opt_times - 1)
                end_index = sample_path_.size() - 1;
            cout << "optimizing trajectory from waypoints index " << start_index
                 << " to index " << end_index << endl;

            vector<Vector3d> input_waypoints(sample_path_.begin() + start_index,
                                             sample_path_.begin() + end_index + 1);
            vector<MatrixXd> part_coeffs;
            vector<double> part_time_seq;
            traj_opt_.setWayPoints(input_waypoints);
            // TODO 更好地选取边界条件
            traj_opt_.setBoundaryCondition(Vector3d(0, 0, 0), Vector3d(0, 0, 0),
                                           Vector3d(0, 0, 0), Vector3d(0, 0, 0));
            if (traj_opt_.optimize() == -1) // -1为nlopt::FAILURE
            {
                ops_ = OPTIMIZE_FAIL;
                ROS_ERROR("Optimizing trajectory failed");
                return;
            }

            traj_opt_.getCoeffs(part_coeffs);
            traj_opt_.getTimeSeq(part_time_seq);

            for (int i = 0; i < 3; i++)
                coeffs_[i].block(start_index, 0, end_index - start_index, 6) = part_coeffs[i];
            time_seq_.insert(time_seq_.end(), part_time_seq.begin(), part_time_seq.end());
        }

        ops_ = OPTIMIZE_TRAJ;

        displayTrajectory(coeffs_, time_seq_, optimized_traj_pub_);
    }

} // namespace Planner